package com.ruoyi.flow.core;

import javax.servlet.*;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletRequestWrapper;
import java.io.IOException;
import java.util.Arrays;
import java.util.List;
import java.util.stream.Collectors;

/**
 * 跨站脚本攻击过滤器
 */
public class XssFilter implements Filter {
    FilterConfig filterConfig = null;

    @Override
    public void init(FilterConfig filterConfig) throws ServletException {
        this.filterConfig = filterConfig;
    }

    @Override
    public void doFilter(ServletRequest servletRequest, ServletResponse servletResponse, FilterChain filterChain) throws IOException, ServletException {
        filterChain.doFilter(new XssHttpServletRequestWrapper((HttpServletRequest) servletRequest), servletResponse);
    }

    @Override
    public void destroy() {
        this.filterConfig = null;
    }

    /**
     * 防止跨站脚本攻击
     * <br/>
     * 所有的请求的getParameter会被替换，如果参数里有敏感词会被替换掉
     */
    class XssHttpServletRequestWrapper extends HttpServletRequestWrapper {
        public XssHttpServletRequestWrapper(HttpServletRequest request) {
            super(request);
        }

        @Override
        public String[] getParameterValues(String paramName) {
            String[] values = super.getParameterValues(paramName);
            if (values == null) {
                return null;
            }
            List<String> encodeValues = Arrays.asList(values).stream().peek(this::cleanXSS).collect(Collectors.toList());
            return encodeValues.toArray(new String[]{});
        }

        @Override
        public String getHeader(String name) {
            String value = super.getHeader(name);
            if (value == null) {
                return null;
            }
            return cleanXSS(value);
        }

        private String cleanXSS(String value) {
            value = value.replaceAll("<", "& lt;").replaceAll(">", "& gt;").replaceAll("\\(", "& #40").replaceAll("\\)", "& #41").replaceAll("'", "& #39").replaceAll("eval\\((.*)\\)", "")
                    .replaceAll("[\\\"\\\'][\\s]*javascript:(.*)[\\\"\\\']", "\"\"");
            return value;
        }
    }
}
